Another same union find question , you can figure out it is a union find problem by the hint that determine a and b connected and a minimal spanning tree kind of a deal.
class Solution {
int[] parent;
int[] size;
public boolean[] distanceLimitedPathsExist(int n, int[][] edgeList, int[][] queries) {
int queriesCount = queries.length;
boolean[] res = new boolean[queriesCount];
// init
parent = new int[n];
size = new int[n];
for(int i = 0; i < n; i ++){
parent[i] = i;
size[i] = 1;
}
// Store original indices with all queries.
int[][] queriesWithIndex = new int[queriesCount][4];
for (int i = 0; i < queriesCount; ++i) {
queriesWithIndex[i][0] = queries[i][0];
queriesWithIndex[i][1] = queries[i][1];
queriesWithIndex[i][2] = queries[i][2];
queriesWithIndex[i][3] = i;
}
Arrays.sort(edgeList,(a,b) ->{return a[2] - b[2];});
Arrays.sort(queriesWithIndex,(a,b) ->{return a[2] - b[2];});
int edgesIndex = 0;
for(int i = 0; i < queriesCount; i ++){
int p = queriesWithIndex[i][0];
int q = queriesWithIndex[i][1];
int limit = queriesWithIndex[i][2];
int originalIndex = queriesWithIndex[i][3];
// because we have the edges sorted
// while our paths are smaller than limit, we want to attache all the edges thta is less than limit
// if this made p q in a same union aka connected, we return true
while(edgesIndex < edgeList.length && edgeList[edgesIndex][2] < limit){
int node1 = edgeList[edgesIndex][0];
int node2 = edgeList[edgesIndex][1];
union(node1,node2);
edgesIndex += 1;
}
res[originalIndex] = find(p) == find(q);
}
return res;
}
private int find(int a){
if(parent[a] == a) return a;
return parent[a] = find(parent[a]);
}
private int union(int a, int b){
int pa = find(a);
int pb = find(b);
if(pa == pb) return 0;
if(size[pa] > pb){
parent[pb] = a;
size[pa] += size[pb];
}else{
parent[pa] = b;
size[pb] += size[pa];
}
return 1;
}
}